import torch
import torch.nn as nn
from torch import Tensor

from models.SENet.SElayer import SELayer
from models.backbone.vgg import vgg16
from utils.conv.adder import adder2d


class BasicBlock(nn.Module):
    def __init__(self, in_channels, features, reduction):
        super().__init__()

        self.conv1 = nn.Conv2d(
            in_channels=in_channels,
            out_channels=features,
            kernel_size=3,
            padding=1,
            bias=False,
        )
        self.se1 = SELayer(features, reduction)
        self.norm1 = nn.BatchNorm2d(num_features=features)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = nn.Conv2d(
            in_channels=features,
            out_channels=features,
            kernel_size=3,
            padding=1,
            bias=False,
        )
        self.se2 = SELayer(features, reduction)
        self.norm2 = nn.BatchNorm2d(num_features=features)
        self.relu2 = nn.ReLU(inplace=True)

    def forward(self, x: Tensor) -> Tensor:
        x = self.conv1(x)
        x = self.se1(x)
        x = self.norm1(x)
        x = self.relu1(x)
        identity = x

        x = self.conv2(x)
        x = self.se2(x)
        x = self.norm2(x)
        x += identity
        x = self.relu2(x)

        return x


class SimSEUnet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, init_features=32, reduction=16):
        super(SimSEUnet, self).__init__()

        features = init_features
        self.encoder1 = BasicBlock(in_channels, features, reduction)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder2 = BasicBlock(features, features * 2, reduction)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder3 = BasicBlock(features * 2, features * 4, reduction)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder4 = BasicBlock(features * 4, features * 8, reduction)
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.bottleneck = BasicBlock(features * 8, features * 16, reduction)

        self.upconv4 = nn.ConvTranspose2d(
            features * 16, features * 8, kernel_size=2, stride=2
        )
        self.decoder4 = BasicBlock((features * 8) * 2, features * 8, reduction)
        self.upconv3 = nn.ConvTranspose2d(
            features * 8, features * 4, kernel_size=2, stride=2
        )
        self.decoder3 = BasicBlock((features * 4) * 2, features * 4, reduction)
        self.upconv2 = nn.ConvTranspose2d(
            features * 4, features * 2, kernel_size=2, stride=2
        )
        self.decoder2 = BasicBlock((features * 2) * 2, features * 2, reduction)
        self.upconv1 = nn.ConvTranspose2d(
            features * 2, features, kernel_size=2, stride=2
        )
        self.decoder1 = BasicBlock(features * 2, features, reduction)

        self.conv = nn.Conv2d(
            in_channels=features, out_channels=out_channels, kernel_size=1
        )

    def forward(self, x):
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))

        bottleneck = self.bottleneck(self.pool4(enc4))

        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.decoder4(dec4)
        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.decoder3(dec3)
        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.decoder2(dec2)
        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.decoder1(dec1)
        return torch.sigmoid(self.conv(dec1))


class BasicAdderBlock(nn.Module):
    def __init__(self, in_channels, features, reduction):
        super(BasicAdderBlock, self).__init__()

        self.conv1 = adder2d(
            input_channel=in_channels,
            output_channel=features,
            kernel_size=3,
            padding=1,
            bias=False,
        )
        self.se1 = SELayer(features, reduction)
        self.norm1 = nn.BatchNorm2d(num_features=features)
        self.relu1 = nn.ReLU(inplace=True)
        self.conv2 = adder2d(
            input_channel=features,
            output_channel=features,
            kernel_size=3,
            padding=1,
            bias=False,
        )
        self.se2 = SELayer(features, reduction)
        self.norm2 = nn.BatchNorm2d(num_features=features)
        self.relu2 = nn.ReLU(inplace=True)

    def forward(self, x: Tensor) -> Tensor:
        x = self.conv1(x)
        x = self.se1(x)
        x = self.norm1(x)
        x = self.relu1(x)
        identity = x

        x = self.conv2(x)
        x = self.se2(x)
        x = self.norm2(x)
        x += identity
        x = self.relu2(x)

        return x


# Using addition instead of multiplication calculation will lead to a large increase in memory usage
class SimAdderSEUnet(nn.Module):
    def __init__(self, in_channels=3, out_channels=1, init_features=32, reduction=16):
        super(SimAdderSEUnet, self).__init__()

        features = init_features
        self.encoder1 = BasicBlock(in_channels, features, reduction)
        self.pool1 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder2 = BasicBlock(features, features * 2, reduction)
        self.pool2 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder3 = BasicBlock(features * 2, features * 4, reduction)
        self.pool3 = nn.MaxPool2d(kernel_size=2, stride=2)
        self.encoder4 = BasicBlock(features * 4, features * 8, reduction)
        self.pool4 = nn.MaxPool2d(kernel_size=2, stride=2)

        self.bottleneck = BasicBlock(features * 8, features * 16)

        self.upconv4 = nn.ConvTranspose2d(
            features * 16, features * 8, kernel_size=2, stride=2
        )
        self.decoder4 = BasicBlock((features * 8) * 2, features * 8, reduction)
        self.upconv3 = nn.ConvTranspose2d(
            features * 8, features * 4, kernel_size=2, stride=2
        )
        self.decoder3 = BasicBlock((features * 4) * 2, features * 4, reduction)
        self.upconv2 = nn.ConvTranspose2d(
            features * 4, features * 2, kernel_size=2, stride=2
        )
        self.decoder2 = BasicBlock((features * 2) * 2, features * 2, reduction)
        self.upconv1 = nn.ConvTranspose2d(
            features * 2, features, kernel_size=2, stride=2
        )
        self.decoder1 = BasicBlock(features * 2, features, reduction)

        self.conv = adder2d(
            input_channel=features, output_channel=out_channels, kernel_size=1
        )

    def forward(self, x):
        enc1 = self.encoder1(x)
        enc2 = self.encoder2(self.pool1(enc1))
        enc3 = self.encoder3(self.pool2(enc2))
        enc4 = self.encoder4(self.pool3(enc3))

        bottleneck = self.bottleneck(self.pool4(enc4))

        dec4 = self.upconv4(bottleneck)
        dec4 = torch.cat((dec4, enc4), dim=1)
        dec4 = self.decoder4(dec4)
        dec3 = self.upconv3(dec4)
        dec3 = torch.cat((dec3, enc3), dim=1)
        dec3 = self.decoder3(dec3)
        dec2 = self.upconv2(dec3)
        dec2 = torch.cat((dec2, enc2), dim=1)
        dec2 = self.decoder2(dec2)
        dec1 = self.upconv1(dec2)
        dec1 = torch.cat((dec1, enc1), dim=1)
        dec1 = self.decoder1(dec1)
        return torch.sigmoid(self.conv(dec1))


class unetUp(nn.Module):
    def __init__(self, in_size, out_size):
        super(unetUp, self).__init__()
        self.conv1 = nn.Conv2d(in_size, out_size, kernel_size=3, padding=1)
        self.conv2 = nn.Conv2d(out_size, out_size, kernel_size=3, padding=1)
        self.up = nn.UpsamplingBilinear2d(scale_factor=2)

    def forward(self, inputs1, inputs2):
        outputs = torch.cat([inputs1, self.up(inputs2)], 1)
        outputs = self.conv1(outputs)
        outputs = self.conv2(outputs)
        return outputs


class SEUnet(nn.Module):
    def __init__(self, num_classes=21, in_channels=3, pretrained=False, model_dir="../model_data"):
        super(SEUnet, self).__init__()
        self.vgg = vgg16(pretrained=pretrained, in_channels=in_channels, model_dir=model_dir)
        in_filters = [192, 384, 768, 1024]
        out_filters = [64, 128, 256, 512]

        self.se1 = SELayer(out_filters[0], 16)
        self.se2 = SELayer(out_filters[1], 16)
        self.se3 = SELayer(out_filters[2], 16)
        self.se4 = SELayer(out_filters[3], 16)
        self.se5 = SELayer(out_filters[3], 16)

        # upsampling
        # 64,64,512
        self.up_concat4 = unetUp(in_filters[3], out_filters[3])
        # 128,128,256
        self.up_concat3 = unetUp(in_filters[2], out_filters[2])
        # 256,256,128
        self.up_concat2 = unetUp(in_filters[1], out_filters[1])
        # 512,512,64
        self.up_concat1 = unetUp(in_filters[0], out_filters[0])

        # final conv (without any concat)
        self.final = nn.Conv2d(out_filters[0], num_classes, 1)

    def forward(self, inputs):
        feat1 = self.vgg.features[:4](inputs)
        feat1 = self.se1(feat1)
        feat2 = self.vgg.features[4:9](feat1)
        feat2 = self.se2(feat2)
        feat3 = self.vgg.features[9:16](feat2)
        feat3 = self.se3(feat3)
        feat4 = self.vgg.features[16:23](feat3)
        feat4 = self.se4(feat4)
        feat5 = self.vgg.features[23:-1](feat4)
        feat5 = self.se5(feat5)

        up4 = self.up_concat4(feat4, feat5)
        up3 = self.up_concat3(feat3, up4)
        up2 = self.up_concat2(feat2, up3)
        up1 = self.up_concat1(feat1, up2)

        final = self.final(up1)

        return final

    def _initialize_weights(self, *stages):
        for modules in stages:
            for module in modules.modules():
                if isinstance(module, nn.Conv2d):
                    nn.init.kaiming_normal_(module.weight)
                    if module.bias is not None:
                        module.bias.data.zero_()
                elif isinstance(module, nn.BatchNorm2d):
                    module.weight.data.fill_(1)
                    module.bias.data.zero_()
